# -*- coding: UTF-8 -*-
import collections
from typing import List
from leetcode_helper import TestcaseHelper, SolutionBase
# Leetcode No: 2127
# https://leetcode-cn.com/problems/maximum-employees-to-be-invited-to-a-meeting/solution/duo-tu-xiang-jie-zhu-shi-tuo-bu-pai-xu-y-lxs3/


class Solution(SolutionBase):
    # 解题的方法，命名和Leetcode的保持一致即可
    def maximumInvitations(self, favorite: List[int]) -> int:
        reversed_graph = collections.defaultdict(list)
        in_deg = collections.defaultdict(int)

        for u, v in enumerate(favorite):
            # relation ship is u --> v
            reversed_graph[v].append(u)
            in_deg[v] += 1

        member_count = len(favorite)
        dp = collections.deque([u for u in range(member_count) if in_deg[u] == 0])

        # 达到剪枝的效果，将环上的树枝去掉
        # 循环结束后，in_deg里入度为0的都是树枝，入度不为0的都是环上的节点
        while dp:
            u = dp.popleft()
            v = favorite[u]
            in_deg[v] -= 1
            if in_deg[v] == 0:
                dp.append(v)

        def bfs_search_circle_depth(root: int):
            nonlocal visited
            depth = 0
            dp = collections.deque([root])

            while dp:
                u = dp.popleft()
                if u in visited:
                    return depth
                else:
                    visited.add(u)
                    dp.append(favorite[u])
                    depth += 1

        ans = 0

        # 计算所有环的长度
        visited = set() # 记录节点是否被访问过
        nodes_in_circles = [u for u in range(member_count) if in_deg[u] > 0] # 保存所有在环上的节点
        two_nodes_circles = [] # 保存长度为2的环上元素
        max_depth_circle = 0 # 最长的环长度

        for node in nodes_in_circles:
            if node in visited:
                continue
            circle_depth = bfs_search_circle_depth(node)
            max_depth_circle = max(max_depth_circle, circle_depth)

            if circle_depth == 2:
                two_nodes_circles.append((node, favorite[node]))

        def bfs_search_2node_circle_depth(root: int):
            dp = collections.deque([root])
            depth = 0
            while dp:
                depth += 1
                for _ in range(len(dp)):
                    u = dp.popleft()
                    for v in reversed_graph[u]:
                        dp.append(v)

            return depth

        # 计算二元环的长度
        max_depth_2node_circle = 0
        if len(two_nodes_circles) == 0:
            ans = max_depth_circle
        else:
            while two_nodes_circles:
                u, v = two_nodes_circles.pop()
                reversed_graph[u] = [i for i in reversed_graph[u] if i != v]
                reversed_graph[v] = [i for i in reversed_graph[v] if i != u]
                left = bfs_search_2node_circle_depth(u)
                right = bfs_search_2node_circle_depth(v)
                max_depth_2node_circle += left + right

        ans = max(max_depth_circle, max_depth_2node_circle)
        return ans


if __name__ == '__main__':
    solution = Solution()

    # 按照leetcode的题目说明，将基本的test case填入
    # 每条case的最后一个元素是正确的结果，即期待的结果
    # 使用元组来保存输入的数据，期待的结果
    testcases = [
        ([3, 0, 1, 4, 1, 2], 4),
        ([2, 2, 1, 2], 3),
        ([1, 2, 0], 3)
    ]

    for case_item in testcases:
        # 获取测试用的case和期待的执行结果
        input1 = case_item[0]
        expect = case_item[len(case_item) - 1]
        exec_result = solution.maximumInvitations(input1)

        # 判断执行结果，输出提示信息
        check_result = solution.check_result(expect, exec_result)
        TestcaseHelper.print_case(check_result, case_item, exec_result)

